{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "代码说明：在此文件中，实现了可见光红外多组对比图生成实验，但是使用的注意力可视化包是魔改的，不是原版，不能在本环境中直接使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 处理代码加载\n",
    "import torch\n",
    "from transformers import AutoProcessor, LlavaForConditionalGeneration,LlavaProcessor\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2\"  # 指定使用第0和第1号GPU\n",
    "\n",
    "# 使用device_map自动分配模型到多个GPU\n",
    "llava_model = LlavaForConditionalGeneration.from_pretrained(\n",
    "    '/data/VLM/llava-v1.5-7b-hf',\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    low_cpu_mem_usage=True,\n",
    "    attn_implementation=\"eager\",\n",
    "    device_map=\"auto\"  # 自动分配到可用GPU\n",
    ")\n",
    "# llava_model = LlavaForConditionalGeneration.from_pretrained('/data/VLM/llava-v1.5-7b-hf', torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, attn_implementation=\"eager\").to('cuda')\n",
    "llava_processor = LlavaProcessor.from_pretrained('/data/VLM/llava-v1.5-7b-hf', patch_size=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# image_path_list=['/data/dataset/M3FD/M3FD_Detection/Vis/00000.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Vis/00385.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Vis/00388.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Vis/00712.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Vis/00716.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Vis/04026.png'\n",
    "#                 ]\n",
    "# lr_image_path_list=[\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Ir/00000.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Ir/00385.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Ir/00388.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Ir/00712.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Ir/00716.png',\n",
    "#                  '/data/dataset/M3FD/M3FD_Detection/Ir/04026.png'\n",
    "#                 ]\n",
    "\n",
    "image_path_list=[\"/data/dataset/sam/images/sa_105.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_109.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_118.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_132.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_134.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_137.jpg\"\n",
    "                 ]\n",
    "lr_image_path_list=[\"/data/dataset/sam/images/sa_105.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_109.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_118.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_132.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_134.jpg\",\n",
    "                 \"/data/dataset/sam/images/sa_137.jpg\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "def normalize_to_255(arr):\n",
    "    # 确保输入是NumPy数组\n",
    "    arr = np.asarray(arr)\n",
    "    # 处理全为常数的情况（避免除以零）\n",
    "    if np.min(arr) == np.max(arr):\n",
    "        return np.full_like(arr, 128, dtype=np.uint8)\n",
    "    # 归一化到0-1范围\n",
    "    normalized_arr = (arr - np.min(arr)) / (np.max(arr) - np.min(arr))\n",
    "    # 缩放至0-255范围\n",
    "    scaled_arr = normalized_arr * 255\n",
    "    # 转换为uint8类型（可选，但通常用于图像数据）\n",
    "    return scaled_arr.astype(np.uint8)\n",
    "\n",
    "def resize_image(arr, pil_img):\n",
    "    # 获取源数组和目标图像的尺寸\n",
    "    src_height, src_width = arr.shape\n",
    "    target_width, target_height = pil_img.size\n",
    "    \n",
    "    # 创建目标数组\n",
    "    target_arr = np.zeros((target_height, target_width), dtype=arr.dtype)\n",
    "    \n",
    "    # 计算缩放因子\n",
    "    x_scale = src_width / target_width\n",
    "    y_scale = src_height / target_height\n",
    "    \n",
    "    # 双线性插值重采样\n",
    "    for y in range(target_height):\n",
    "        for x in range(target_width):\n",
    "            # 计算在源数组中的对应坐标\n",
    "            src_x = x * x_scale\n",
    "            src_y = y * y_scale\n",
    "            # 找到周围的四个点\n",
    "            x0 = int(src_x)\n",
    "            y0 = int(src_y)\n",
    "            x1 = min(x0 + 1, src_width - 1)\n",
    "            y1 = min(y0 + 1, src_height - 1)\n",
    "            # 计算权重\n",
    "            fx = src_x - x0\n",
    "            fy = src_y - y0\n",
    "            # 双线性插值\n",
    "            value = (1 - fx) * (1 - fy) * arr[y0, x0] + \\\n",
    "                    fx * (1 - fy) * arr[y0, x1] + \\\n",
    "                    (1 - fx) * fy * arr[y1, x0] + \\\n",
    "                    fx * fy * arr[y1, x1]\n",
    "            # 填充目标数组\n",
    "            target_arr[y, x] = value\n",
    "    return target_arr\n",
    "\n",
    "def stretch_array_to_image(arr, pil_img, composite=False,color=1):\n",
    "    target_width, target_height = pil_img.size\n",
    "    target_arr = resize_image(arr, pil_img)\n",
    "\n",
    "    if composite:\n",
    "        pil_img = pil_img.convert('RGBA')\n",
    "        # 将二值图转换为RGBA并指定颜色（例如红色）\n",
    "        binary_rgba = Image.new('RGBA', pil_img.size)\n",
    "        pixels = binary_rgba.load()\n",
    "        # 转换为PIL图像\n",
    "        for y in range(target_width):\n",
    "            for x in range(target_height):\n",
    "                if color==1:\n",
    "                    pixels[y, x] = (target_arr[x, y], 0, 0, 255)\n",
    "                elif color==2:\n",
    "                    pixels[y, x] = (0, target_arr[x, y], target_arr[x, y], 255)\n",
    "                \n",
    "        result = Image.alpha_composite(pil_img, binary_rgba)\n",
    "        return result\n",
    "    else:\n",
    "        return Image.fromarray(target_arr)\n",
    "\n",
    "\n",
    "\n",
    "def contract_att_map(vis_att_map,lr_att_map,image):\n",
    "    target_width, target_height = image.size\n",
    "    vis_target_arr = resize_image(vis_att_map, image)\n",
    "    lr_target_arr = resize_image(lr_att_map, image)\n",
    "    # 转换为PIL图像\n",
    "    pil_img = image.convert('RGBA')\n",
    "        # 将二值图转换为RGBA并指定颜色（例如红色）\n",
    "    binary_rgba = Image.new('RGBA', pil_img.size)\n",
    "    pixels = binary_rgba.load()\n",
    "    for y in range(target_width):\n",
    "        for x in range(target_height):\n",
    "            pixels[y, x] = (vis_target_arr[x, y], lr_target_arr[x, y], lr_target_arr[x, y], 150)\n",
    "                \n",
    "    result = Image.alpha_composite(pil_img, binary_rgba)\n",
    "    return result\n",
    "\n",
    "from attention_utils.llava_methods import *\n",
    "\n",
    "def generate_attention_map(model, processor, image_path, lr_image_path, layer, head, att_type):\n",
    "    image = Image.open(image_path).convert(\"RGB\")\n",
    "    lr_image = Image.open(lr_image_path).convert(\"RGB\")\n",
    "\n",
    "    question = 'Find all the people in the picture and locate their positions.'\n",
    "    general_question = 'Write a general description of the image.'\n",
    "\n",
    "    prompt = f\"<image>\\nUSER: {question}\\nASSISTANT:\"\n",
    "    general_prompt = f\"<image>\\nUSER: {general_question}\\nASSISTANT:\"\n",
    "    if att_type==\"head_rel\":\n",
    "        vis_att_map = head_rel_attention_llava(image, prompt, general_prompt, model, processor, layer, head)\n",
    "        lr_att_map = head_rel_attention_llava(lr_image, prompt, general_prompt, model, processor, layer, head)\n",
    "    elif att_type==\"head_orin\":\n",
    "        vis_att_map = head_orin_attention_llava(image, prompt, model, processor, layer, head)\n",
    "        lr_att_map = head_orin_attention_llava(lr_image, prompt, model, processor, layer, head)\n",
    "    elif att_type==\"rel\":\n",
    "        vis_att_map = rel_attention_llava(image, prompt, general_prompt, model, processor, layer)\n",
    "        lr_att_map = rel_attention_llava(lr_image, prompt, general_prompt, model, processor, layer)\n",
    "    elif att_type==\"grad\":\n",
    "        vis_att_map=gradient_attention_llava(image, prompt, general_prompt, model, processor, layer)\n",
    "        lr_att_map=gradient_attention_llava(lr_image, prompt, general_prompt, model, processor, layer)\n",
    "    elif att_type==\"orin\":\n",
    "        vis_att_map=orin_attention_llava(image, prompt, model, processor, layer)\n",
    "        lr_att_map=orin_attention_llava(lr_image, prompt, model, processor, layer)\n",
    "    elif att_type==\"pure\":\n",
    "        vis_att_map=pure_gradient_llava(image, prompt, general_prompt, model, processor)\n",
    "        lr_att_map=pure_gradient_llava(lr_image, prompt, general_prompt, model, processor)\n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "    vis_att_map=normalize_to_255(vis_att_map)\n",
    "    lr_att_map=normalize_to_255(lr_att_map)\n",
    "    # result_map1=stretch_array_to_image(vis_att_map, image ,composite=True,color=1)\n",
    "    # result_map2=stretch_array_to_image(lr_att_map, image ,composite=True,color=2)\n",
    "    att_map1=contract_att_map(vis_att_map,lr_att_map,image)\n",
    "    att_map2=contract_att_map(vis_att_map,lr_att_map,lr_image)\n",
    "    return [image, lr_image, att_map1, att_map2]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_pack=[]\n",
    "for i in range(len(image_path_list)):\n",
    "    image_pack.append(generate_attention_map(llava_model, llava_processor, image_path_list[i], lr_image_path_list[i],layer=14,head=24,att_type=\"pure\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "def visualize_image_grid(images, save_path=None):\n",
    "    \"\"\"\n",
    "    Visualize a list of PIL images in a grid layout.\n",
    "    \n",
    "    Args:\n",
    "        images: List of lists of PIL images with shape (6, 4, 768, 1024)\n",
    "        save_path: Optional path to save the visualization\n",
    "    \"\"\"\n",
    "    # Create a figure with 6 rows and 4 columns\n",
    "    fig, axes = plt.subplots(6, 4, figsize=(20, 30))\n",
    "    \n",
    "    # Plot each image\n",
    "    for i in range(6):  # 6 rows\n",
    "        for j in range(4):  # 4 columns\n",
    "            if i < len(images) and j < len(images[i]):\n",
    "                # Convert PIL image to numpy array\n",
    "                img_array = np.array(images[i][j])\n",
    "                axes[i, j].imshow(img_array)\n",
    "                axes[i, j].axis('off')  # Turn off axis\n",
    "    \n",
    "    # Adjust layout to prevent overlap\n",
    "    plt.tight_layout()\n",
    "    \n",
    "    # Save the figure if save_path is provided\n",
    "    if save_path:\n",
    "        plt.savefig(save_path, bbox_inches='tight', dpi=300)\n",
    "    \n",
    "    plt.show()\n",
    "# Example usage:\n",
    "# images = [PIL.Image.open(path) for path in image_paths]\n",
    "# visualize_image_grid(images, save_path='output.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 假设你的图片列表是 images\n",
    "visualize_image_grid(image_pack, save_path='output.png')  # 可选是否保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llava",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
